
#pragma once
#include <gtest/gtest.h>
#include <spatialindex/SpatialIndex.h>
#include <spdlog/spdlog.h>
#include <algorithm>
#include <cmath>
#include <entt/entt.hpp>
#include <iostream>
#include <random>
#include <vector>

// using namespace SpatialIndex;
using namespace SpatialIndex;
namespace spi = SpatialIndex;

// struct Point : spi::Point {
//   Point(spi::Point& p) : spi::Point(p, 2) {}
//   double x() { return m_pCoords[0]; }
//   double y() { return m_pCoords[1]; }
// };

// struct Box {
//   double x, y;
//   double width, height;

//   Box(double x, double y, double width, double height) : x(x), y(y), width(width), height(height) {}

//   void center() {
//     x = (x + width / 2.0);
//     y = (y + height / 2.0);
//     return Point(x, y);
//   }
// };

// struct Box : spi::Region {
//   Box(const Point& low, const Point& high) : Region(p1, p2){};

//   double x() { return m_pLow[0]; }
//   double y() { return m_pLow[1]; }
//   double width() { return m_pLow[1] - m_pLow[0]; }
//   double height() { return m_pHigh[1] - m_pHigh[0]; }
//   double cx() { return (m_pLow[1] - m_pLow[0]) / 2.0; }
//   double cy() { return (m_pHigh[1] - m_pHigh[0]) / 2.0; }

//   Point center() { return Point(x + width() / 2.0, y + height() / 2.0); }
// };

class MyVisitor : public spi::IVisitor {
 public:
  size_t m_indexIO{0};
  size_t m_leafIO{0};
  std::vector<size_t> ids;

 public:
  MyVisitor() = default;

  void visitNode(const INode& n) override {
    if (n.isLeaf())
      m_leafIO++;
    else
      m_indexIO++;
  }

  void visitData(const IData& d) override { ids.push_back(d.getIdentifier()); }

  void visitData(std::vector<const IData*>& v) override { std::cout << v[0]->getIdentifier() << " " << v[1]->getIdentifier() << std::endl; }
};

struct LayoutBlock {
  std::vector<entt::entity> objects;
  SpatialIndex::Region mbr;
  uint64_t hilbert_index;
  std::string chunk_id;

  void compute_mbr(const entt::registry& reg) {
    if (objects.empty())
      return;

    auto& mbr_0 = reg.get<Region>(objects[0]);
    double low[2] = {mbr_0.m_pLow[0], mbr_0.m_pLow[1]};
    double high[2] = {mbr_0.m_pHigh[0], mbr_0.m_pHigh[1]};
    for (const auto& obj : objects) {
      auto& mbr = reg.get<Region>(obj);
      for (int i = 0; i < 2; ++i) {
        low[i] = std::min(low[i], mbr.m_pLow[i]);
        high[i] = std::max(high[i], mbr.m_pHigh[i]);
      }
    }
    mbr = SpatialIndex::Region(low, high, 2);
  }
};

struct GeometryObject {};

static void rot(int n, uint32_t& x, uint32_t& y, int rx, int ry) {
  if (ry == 0) {
    if (rx == 1) {
      x = n - 1 - x;
      y = n - 1 - y;
    }
    std::swap(x, y);
  }
}

uint64_t xy_to_hilbert(uint32_t x, uint32_t y, int bits) {
  uint64_t d = 0;
  int n = 1 << bits;
  for (int s = n / 2; s > 0; s /= 2) {
    int rx = (x & s) ? 1 : 0;
    int ry = (y & s) ? 1 : 0;
    d += (uint64_t)s * s * ((3 * rx) ^ ry);
    rot(s, x, y, rx, ry);
  }
  return d;
}

// 按 Hilbert 排序
void sort_by_hilbert(const entt::registry& reg, std::vector<entt::entity>& objs, const Region& bbox, int bits = 16) {
  double x_min = bbox.m_pLow[0], y_min = bbox.m_pLow[1];
  double range = std::max(bbox.m_pHigh[0] - x_min, bbox.m_pHigh[1] - y_min);
  uint32_t scale = (1U << bits) - 1;

  std::sort(objs.begin(), objs.end(), [&](const entt::entity& a, const entt::entity& b) {
    auto& mbr_a = reg.get<Region>(a);
    auto& mbr_b = reg.get<Region>(b);

    auto cx_a = (mbr_a.m_pLow[0] + mbr_a.m_pHigh[0]) / 2.0;
    auto cx_b = (mbr_b.m_pLow[0] + mbr_b.m_pHigh[0]) / 2.0;
    auto cy_a = (mbr_a.m_pLow[1] + mbr_a.m_pHigh[1]) / 2.0;
    auto cy_b = (mbr_b.m_pLow[1] + mbr_b.m_pHigh[1]) / 2.0;

    uint32_t nx_a = static_cast<uint32_t>((cx_a - x_min) / range * scale);
    uint32_t ny_a = static_cast<uint32_t>((cy_a - y_min) / range * scale);
    uint32_t nx_b = static_cast<uint32_t>((cx_b - x_min) / range * scale);
    uint32_t ny_b = static_cast<uint32_t>((cy_b - y_min) / range * scale);

    return xy_to_hilbert(nx_a, ny_a, bits) < xy_to_hilbert(nx_b, ny_b, bits);
  });
}

std::vector<LayoutBlock> str_partition(const entt::registry& reg, std::vector<entt::entity>& objects, size_t leaf_capacity,
                                       const SpatialIndex::Region& design_bbox) {

  if (objects.empty())
    return {};

  // Step 1: Hilbert 排序
  sort_by_hilbert(reg, objects, design_bbox);

  // Step 2: 按 leaf_capacity 切分（STR 的简化版）
  std::vector<LayoutBlock> blocks;
  size_t n = objects.size();
  size_t num_blocks = (n + leaf_capacity - 1) / leaf_capacity;

  for (size_t i = 0; i < num_blocks; ++i) {
    size_t start = i * leaf_capacity;
    size_t end = std::min(start + leaf_capacity, n);
    LayoutBlock block;
    block.objects.assign(objects.begin() + start, objects.begin() + end);
    block.compute_mbr(reg);
    blocks.push_back(block);
  }

  return blocks;
}

// 构建全局查询 R 树（用于运行时）
void build_global_rtree(const std::vector<LayoutBlock>& blocks, const std::string& rtree_path) {
  // 使用 libspatialindex 创建磁盘 R 树
  //   SpatialIndex::IStorageManager* storage = SpatialIndex::StorageManager::createNewDiskStorageManager(rtree_path, 4096);
  //   SpatialIndex::ISpatialIndex* tree = SpatialIndex::RTree::createNewRTree(*storage, 0.7, 100, 100, 2, SpatialIndex::RTree::RV_RSTAR);

  std::string baseName = rtree_path;
  IStorageManager* diskfile = StorageManager::createNewDiskStorageManager(baseName, 4096);
  StorageManager::IBuffer* file = StorageManager::createNewRandomEvictionsBuffer(*diskfile, 10, false);

  id_type indexIdentifier;
  ISpatialIndex* tree = RTree::createNewRTree(*file, 0.7, 100, 10, 2, SpatialIndex::RTree::RV_RSTAR, indexIdentifier);

  for (size_t i = 0; i < blocks.size(); ++i) {
    // 将 chunk_id 作为 data
    std::string chunk_id = "chunk_" + std::to_string(i);
    size_t len = chunk_id.size();
    uint8_t* data = new uint8_t[len];
    std::copy(chunk_id.begin(), chunk_id.end(), data);

    tree->insertData(len, data, blocks[i].mbr, i);
    delete[] data;
  }

  delete tree;
  delete file;
  delete diskfile;
}

TEST(AggregatedTreeTest, BasicAssertion) {
  // 1. 模拟加载几何对象（实际从 GDS/OASIS 解析）
  std::vector<entt::entity> geometries;
  entt::registry reg;

  std::random_device rd;
  std::mt19937 gen(rd());
  std::uniform_real_distribution<double> dis(0.0, 100.0);  // [0,1)
                                                           //   double r = dis(gen);
                                                           //   std::cout << "随机浮点数: " << r << std::endl;

  for (uint32_t i = 0; i < 1e2; ++i) {
    auto e = reg.create();
    // double x1 = dis(gen);
    // double y2 = dis(gen);
    double low[2], high[2];
    // low[0] = dis(gen);
    // low[1] = dis(gen);
    // high[0] = dis(gen);
    // high[1] = dis(gen);

    low[0] = i * 1.0;
    low[1] = i * 1.0;
    high[0] = i + 1.0;
    high[1] = i + 1.0;
    SpatialIndex::Region r(low, high, 2);
    reg.emplace<Region>(e, r);

    geometries.push_back(e);
  }
  // ... 填充 geometries

  // 2. 设计边界
  double low[2] = {0.0, 0.0};
  double high[2] = {100.0, 100.0};  // 1cm x 1cm
  SpatialIndex::Region design_bbox(low, high, 2);

  // 3. STR + Hilbert 分块
  size_t LEAF_CAPACITY = 50;  // 每区块最多50万个对象
  auto blocks = str_partition(reg, geometries, LEAF_CAPACITY, design_bbox);

  // 4. 为每个区块计算 Hilbert 索引（用于排序）
  double x_min = design_bbox.m_pLow[0], y_min = design_bbox.m_pLow[1];
  double range = std::max(design_bbox.m_pHigh[0] - x_min, design_bbox.m_pHigh[1] - y_min);
  const int BITS = 16;
  uint32_t scale = (1U << BITS) - 1;

  for (auto& block : blocks) {
    double cx = (block.mbr.m_pLow[0] + block.mbr.m_pHigh[0]) / 2.0;
    double cy = (block.mbr.m_pLow[1] + block.mbr.m_pHigh[1]) / 2.0;
    uint32_t nx = static_cast<uint32_t>((cx - x_min) / range * scale);
    uint32_t ny = static_cast<uint32_t>((cy - y_min) / range * scale);
    block.hilbert_index = xy_to_hilbert(nx, ny, BITS);
  }

  // 5. 按 Hilbert 排序
  std::sort(blocks.begin(), blocks.end(), [](const LayoutBlock& a, const LayoutBlock& b) { return a.hilbert_index < b.hilbert_index; });

  // 6. 序列化存储
  //   std::string chunk_dir = "chunks/";
  //   // system("mkdir -p chunks");

  //   for (size_t i = 0; i < blocks.size(); ++i) {
  //     std::ostringstream oss;
  //     oss << chunk_dir << "chunk_" << std::setfill('0') << std::setw(8) << i << ".bin";
  //     serialize_block(blocks[i], oss.str());
  //     blocks[i].chunk_id = "chunk_" + std::to_string(i);
  //   }

  // 7. 构建全局 R 树索引
  build_global_rtree(blocks, "rtree_test");

  // 8. 查询 R 树
  std::string baseName = "rtree_test";
  IStorageManager* diskfile = StorageManager::loadDiskStorageManager(baseName);
  StorageManager::IBuffer* file = StorageManager::createNewRandomEvictionsBuffer(*diskfile, 10, false);
  ISpatialIndex* tree = RTree::loadRTree(*file, 1);

  // 3. 范围查询：查找与 [1.5, 3.5] x [1.5, 3.5] 相交的对象
  double p1[2], p2[2];
  p1[0] = 1.5;
  p1[1] = 1.5;
  p2[0] = 3.5;
  p2[1] = 3.5;
  Region queryRegion(Point(p1, 2), Point(p2, 2));

  // MyVisitor visitor;
  MyVisitor visitor;
  tree->intersectsWithQuery(queryRegion, visitor);

  std::cout << "Range query results (IDs): ";
  for (auto& id : visitor.ids) {
    std::cout << id << " ";
  }
  std::cout << std::endl;

  //   // 4. 最近邻查询：查找离点 (0.5, 0.5) 最近的 2 个对象
  //   double p1[2];
  //   p1[0] = 0.5;
  //   p1[1] = 0.5;
  //   Point queryPoint(p1, 2);
  //   //   MyVisitor visitor;
  //   MyVisitor knnVisitor;
  //   std::vector<spi::id_type> knnResults;
  //   tree->nearestNeighborQuery(2, queryPoint, knnVisitor);

  //   std::cout << "kNN query results (IDs): ";
  //   for (auto& id : knnVisitor.ids) {
  //     std::cout << id << " ";
  //   }
  //   std::cout << std::endl;

  std::cout << "Partitioned into " << blocks.size() << " blocks.\n";
}